'''
数据集分割, 训练、验证比例为8:2
'''
import os
import PIL.Image as Image
import cv2
import numpy as np
import torch.utils.data as data
import torchvision.transforms as transforms
from utils.config import *

# root = 'G:/KITTI/KITTI2015/data_scene_flow' # Windows
# root = '/media/ubuntu/e/zhouyiqing/KITTI/KITTI2015/data_scene_flow' # Ubuntu

# KITTI2015数据集
class KITTI2015(data.Dataset):
    def __init__(self, root, transform, mode, occ=True):
        self.root = root
        self.transform = transform
        self.mode = mode

        validate_rate = 0.2 # 验证集占比

        # 拼接左右图像文件夹路径
        if mode == 'train' or mode == 'validate':
            root = os.path.join(root, 'training')
        elif mode == 'test':
            root = os.path.join(root, 'testing')
        left_root = os.path.join(root, 'image_2')
        right_root = os.path.join(root, 'image_3')

        # 对应图片在文件夹的范围
        if mode == 'train':
            imgs_range = range(int(200 * (1 - validate_rate )))
        elif mode == 'validate':
            imgs_range = range(int(200 * (1 - validate_rate)), 200)
        elif mode == 'test':
            imgs_range = range(200)

        img_fmt = '{:06}_10.png'
        left_imgs = list()
        right_imgs = list()

        # 拼接图片文件名，并加入文件名list
        for i in imgs_range:
            left_imgs.append(os.path.join(left_root, img_fmt.format(i)))
            right_imgs.append(os.path.join(right_root, img_fmt.format(i)))

        self.left_imgs = left_imgs
        self.right_imgs = right_imgs

        # 训练集和测试集含有视差文件
        if mode == 'train' or mode == 'validate':
            disp_fmt = '{:06}_10.png'
            disp_imgs = list()

            # 区分有无遮挡，occ为所有区域(排行榜中-all)
            if occ:
                disp_root = os.path.join(root, 'disp_occ_0')
            else:
                disp_root = os.path.join(root, 'disp_noc_0')

            # 拼接视差文件名，并加入文件名list
            for i in imgs_range:
                disp_imgs.append(os.path.join(disp_root, disp_fmt.format(i)))

            self.disp_imgs = disp_imgs

    def __getitem__(self, index):
        # 读取RGB图像
        left_img = Image.open(self.left_imgs[index]).convert('RGB')
        right_img = Image.open(self.right_imgs[index]).convert('RGB')
        if system == 'Ubuntu': img_name = self.left_imgs[index].split('/')[-1]
        else: img_name = self.left_imgs[index].split('\\')[-1]

        if self.mode != 'test':
            # 读取视差灰度图像
            disp_img = Image.open(self.disp_imgs[index])

        # 图像预处理
        if self.transform:
            seed = torch.random.seed()
            torch.random.manual_seed(seed)
            left_img = self.transform(left_img)
            torch.random.manual_seed(seed)
            right_img = self.transform(right_img)

            # 数据保存到字典
            if self.mode != 'test':
                disp_transform = transforms.Compose([transforms.RandomCrop((height, width))]) # 随机裁剪
                torch.random.manual_seed(seed)
                disp_img = np.array(disp_transform(disp_img), dtype=np.float32) / 256
                if under_sampling:
                    disp_img = disp_img // 2

                # 张量转图片显示
                # transforms.ToPILImage()(left_img).convert('RGB').show()
                # transforms.ToPILImage()(right_img).convert('RGB').show()
                # transforms.ToPILImage()(disp_img).convert('L').show()

                data = {'left': left_img, 'right': right_img, 'disp': disp_img, 'name': img_name}
            else:
                data = {'left': left_img, 'right': right_img, 'name': img_name}

        return data

    def __len__(self):
        return len(self.left_imgs)
